
#include "aggregation.hpp"

#include "csrMatrix.hpp"
#include "pmisAggregation.hpp"

namespace UNAP {

template <typename Cmpt>
Aggregation<Cmpt>::Aggregation(const CSRMatrix<Cmpt> &A,
                               label rowSizeInCoarsestLevel)
    : Coarsening<Cmpt>(A.comm()), finestMatrix_(&A), alpha_(1.0) {
  this->rowSizeInCoarsestLevel_ = rowSizeInCoarsestLevel;
  SET_maxLevels(this->maxLevels_);
}

template <typename Cmpt>
const Matrix<Cmpt> &Aggregation<Cmpt>::matrixLevel(const label i) const {
  if (i == 0) {
    return *finestMatrix_;
  } else {
    return coarseMatrixLevels_[i - 1];
  }
}
template <typename Cmpt>
const Matrix<Cmpt> &Aggregation<Cmpt>::coarseMatrix(const label i) const {
  return coarseMatrixLevels_[i];
}

template <typename Cmpt>
const CSRMatrix<Cmpt> &Aggregation<Cmpt>::Pmatrix(const label leveli) const {
  return P_[leveli];
}

template <typename Cmpt>
const CSRMatrix<Cmpt> &Aggregation<Cmpt>::Rmatrix(const label leveli) const {
  return R_[leveli];
}

template <typename Cmpt>
void Aggregation<Cmpt>::agglomerate() {
  this->nCreatedLevels_ = 0;
  CSRMatrix<Cmpt> *A = const_cast<CSRMatrix<Cmpt> *>(finestMatrix_);
  if (A->tsize() <= this->rowSizeInCoarsestLevel_) {
    this->noCoarse_ = true;
    return;
  }

  while (this->nCreatedLevels_ < this->maxLevels_ - 1) {
    if (A->tsize() < this->rowSizeInCoarsestLevel_) {
      break;
    } else {
      if (!agglomerate(*A, this->nCreatedLevels_)) break;
      A = &coarseMatrixLevels_[this->nCreatedLevels_];
      this->nCreatedLevels_++;
    }
  }

  if (this->nCreatedLevels_) {
    compactLevels(this->nCreatedLevels_);
    this->noCoarse_ = false;
  }

#ifdef DEBUG
  this->commcator_->log() << this->nCreatedLevels_ << " coarse levels created!"
                          << ENDL;
#endif
}

template <typename Cmpt>
bool Aggregation<Cmpt>::agglomerate(const CSRMatrix<Cmpt> &fineA, label level) {
  PMISAggregation<Cmpt> *pmis;
  pmis = new PMISAggregation<Cmpt>(fineA);
  label gSize;
  const label lSize = pmis->size();
  this->commcator_->allReduce("sumLocalNewNnz", &lSize, &gSize, 1, COMM_LABEL,
                              COMM_SUM);
  this->commcator_->finishTask("sumLocalNewNnz");
  if (gSize < this->rowSizeInCoarsestLevel_) return false;

  CSRMatrix<Cmpt> *P = new CSRMatrix<Cmpt>(pmis->P());
  CSRMatrix<Cmpt> *R = P->transpose();
  P_.setLevel(level, *P, false);
  R_.setLevel(level, *R, false);
  CSRMatrix<Cmpt> *tmp = product(fineA, *P);
  CSRMatrix<Cmpt> *Ac = product(*R, *tmp);
  Ac->scale(1. / alpha_);
  Ac->splitOffDiag();
  coarseMatrixLevels_.setLevel(level, *Ac, false);
  delete tmp;
  delete pmis;
  return true;
}

template <typename Cmpt>
void Aggregation<Cmpt>::SET_maxLevels(const label i) {
  this->rowSize_.SET_size(i);
  coarseMatrixLevels_.SET_size(i);
  this->maxLevels_ = i;
}

template <typename Cmpt>
void Aggregation<Cmpt>::compactLevels(const label nCreatedLevels) {
  this->nCreatedLevels_ = nCreatedLevels;
  this->rowSize_.SET_size(nCreatedLevels);
  coarseMatrixLevels_.SET_size(nCreatedLevels);
  P_.SET_size(nCreatedLevels);
  R_.SET_size(nCreatedLevels);
}

template class Aggregation<scalar>;
}  // namespace UNAP
